# -*- coding: utf-8 -*-

import paramiko
import time
import os
import logging
import sys
from datetime import datetime
from dingtalkchatbot.chatbot import DingtalkChatbot


current_index = 0  # 初始索引为0

# 配置日志记录
logging.basicConfig(
    filename="transfer.log",
    level=logging.ERROR,
    format="%(asctime)s - %(levelname)s - %(message)s",
)

# 定义源服务器和目标服务器的信息
source_server = {"host": "182.254.192.210", "port": "22"}

# 日志文件和目标目录的映射
log_dirs = {
    "/tmp/img_path.log": "/data/images",
}

# 重试次数和等待时间
max_retries = 3  # 最大重试次数
retry_interval = 300  # 5分钟等待时间，单位是秒


def nowTime():
    now = datetime.now()
    # 将时间格式化为字符串
    formatted_time = now.strftime("%Y-%m-%d %H:%M:%S")
    return formatted_time


# 发送钉钉通知的函数
def send_dingtalk_alert(message):
    # 设置 Webhook 地址
    url = "https://oapi.dingtalk.com/robot/send?access_token=c54fa6837d584b0abc56c0a7d0456ca37d59213dcc7ea3f6f104b20de3c71284"
    secret = "SEC9de9f15e7d9bc3be3c0c5f50adcf283ff26189940919486b9cdb0feda9ee45b1"

    # 初始化机器人
    bot = DingtalkChatbot(webhook=url, secret=secret)

    try:
        # 发送请求
        bot.send_text(message)
        print(message)
    except Exception as e:
        error_message = f"发送钉钉通知时发生异常：{e}"
        print(error_message)
        logging.error(error_message)


# 初始化SSH客户端连接函数（包括隧道重连）
def initialize_ssh_client(server_info):
    retries = 0
    client = paramiko.SSHClient()
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())

    while retries < max_retries:
        try:
            # 使用SSH密钥认证，不需要用户名和密码
            client.connect(
                server_info["host"],
                port=int(server_info["port"]),
                key_filename="/root/.ssh/id_rsa",
            )
            break
        except paramiko.ssh_exception.SSHException as e:
            error_message = f"SSH 隧道连接失败: {e}"
            print(error_message)
            logging.error(error_message)
            retries += 1
            if retries < max_retries:
                print("5分钟后尝试重新连接...")
                time.sleep(retry_interval)  # 休眠5分钟后重试
            else:
                current_time = nowTime()
                print("无法连接，发送钉钉通知...")
                send_dingtalk_alert(
                    f"*{current_time}，SSH连接失败，无法恢复，请关注！*"
                )

    return client


# 传输文件（包括断开连接时的续传）
def transfer_file(local_file, target_dir, ssh_client, start_index):
    sftp = ssh_client.open_sftp()
    retries = 0  # 添加重试次数计数器

    for i in range(start_index, len(local_file)):
        file_path = local_file[i]
        source_file = os.path.basename(file_path)
        target_file = os.path.join(target_dir, source_file)

        while retries < max_retries:  # 在连接失败时重试
            try:
                # 传输文件
                sftp.put(file_path, target_file)
                print(f"正在传输文件: {source_file}")

                # 保存成功传输的文件索引
                save_last_index(i)
                retries = 0  # 传输成功，重置重试计数

                start_index = current_index
                print(f"start_index={start_index}")
                index = start_index + 1
                print(f"index={index}")
                print(f"len(local_file)={len(local_file)}")

                # 如果已经成功传输了所有文件，则直接跳过
                if index >= len(local_file):
                    current_time = nowTime()
                    access_message = f"*{current_time}，所有文件已成功传输，请关注！*"
                    print(access_message)
                    send_dingtalk_alert(access_message)
                    # 关闭SFTP连接
                    sftp.close()
                    sys.exit(0)  # 退出程序

                break  # 传输成功，退出重试循环

            except paramiko.ssh_exception.SSHException as e:
                current_time = nowTime()
                print(f"{current_time}，{source_file} 传输失败！")
                error_message = f"SSH 隧道连接失败: {e}"
                print(error_message)
                logging.error(error_message)
                retries += 1
                # 记录传输进度
                save_last_index(i)
                if retries < max_retries:
                    print(
                        f"{current_time}，SSH连接失败，等待 {retry_interval} 秒后尝试重新连接..."
                    )
                    time.sleep(retry_interval)  # 休眠5分钟后重试
                    # 关闭SFTP连接
                    sftp.close()
                    # 重新初始化SSH连接
                    ssh_client = initialize_ssh_client(source_server)
                    sftp = ssh_client.open_sftp()
                else:
                    current_time = nowTime()
                    print("无法连接，发送钉钉通知...")
                    access_message = (
                        f"*{current_time}，SSH连接失败，无法恢复，请关注！*"
                    )
                    send_dingtalk_alert(access_message)
                    return False  # 表示文件传输失败

            except Exception as e:
                current_time = nowTime()
                error_message = f"文件 {source_file} 传输错误: {e}"
                print(error_message)
                logging.error(error_message)
                access_message = (
                    f"*{current_time}，文件 {source_file} 传输失败，请关注！*"
                )
                send_dingtalk_alert(access_message)
                # 记录传输进度
                save_last_index(i)
                return False  # 表示文件传输失败

        if retries >= max_retries:
            return False  # 表示文件传输失败

    sftp.close()
    return True  # 返回True表示文件传输成功


# 读取日志文件
def read_log_file(log_file):
    with open(log_file, "r") as file:
        content = file.read()
    return content.splitlines()  # 使用 splitlines() 方法拆分为行列表


# 保存成功传输的文件索引
def save_last_index(index):
    global current_index
    current_index = index


def read_last_index():
    global current_index
    return current_index


def main():
    # 初始化SSH客户端连接
    source_ssh = initialize_ssh_client(source_server)

    # 遍历每个日志文件和目标目录的映射
    for local_log_file, target_dir in log_dirs.items():
        # 获取上次成功传输的文件索引
        start_index = read_last_index()

        retries = 0
        all_files_transferred = False  # 标志变量，跟踪是否所有文件都已传输成功

        try:
            # 读取需要传输的文件列表
            with open(local_log_file, "r") as file:
                files_to_transfer = file.read().splitlines()

            # 如果已经成功传输了所有文件，则直接跳过
            if start_index >= len(files_to_transfer):
                print(len(files_to_transfer))
                current_time = nowTime()
                access_message = f"*{current_time}，所有文件已成功传输，请关注！*"
                source_ssh.close()
                send_dingtalk_alert(access_message)
                break

            while retries < max_retries:
                # 检查SSH连接是否仍然有效，如果无效则重新连接
                if not source_ssh.get_transport().is_active():
                    source_ssh.close()
                    source_ssh = initialize_ssh_client(source_server)

                transfer_result = transfer_file(
                    files_to_transfer,
                    target_dir,
                    source_ssh,
                    start_index,
                )

                if not transfer_result:
                    retries += 1  # 如果需要重新连接SSH，增加重试计数
                    if retries < max_retries:
                        print(f"SSH连接失败，等待 {retry_interval} 秒后尝试重新连接...")
                        time.sleep(retry_interval)  # 等待5分钟后重试
                        # 重新初始化SSH连接
                        source_ssh.close()
                        source_ssh = initialize_ssh_client(source_server)
                    else:
                        current_time = nowTime()
                        print("无法连接，发送钉钉通知...")
                        access_message = (
                            f"*{current_time}，SSH连接失败，无法恢复，请关注！*"
                        )
                        send_dingtalk_alert(access_message)
                        break  # 退出内层while循环
                else:
                    # 如果传输成功，更新内存中的索引
                    save_last_index(start_index)  # 保存索引到内存中

                    # 检查是否所有文件都已成功传输
                    if start_index >= len(files_to_transfer):
                        all_files_transferred = True
                        break  # 退出内层while循环

            if all_files_transferred:
                current_time = nowTime()
                access_message = f"*{current_time}，所有文件已成功传输，请关注！*"
                source_ssh.close()
                print(access_message)
                send_dingtalk_alert(access_message)
                sys.exit(0)  # 退出程序

        except Exception as e:
            print(f"遇到错误: {e}")
            retries += 1
            if retries < max_retries:
                print(f"等待 {retry_interval} 秒后继续...")
                time.sleep(retry_interval)  # 休眠5分钟
            continue

        finally:
            source_ssh.close()


if __name__ == "__main__":
    main()
